package com.ai.bayesi.core;

import java.util.*;

/**
 * 贝叶斯网络节点类
 */
public class Node {
    private String name;
    private List<String> values;
    private List<Node> parents;
    private Map<List<String>, double[]> cpt; // 条件概率表

    public Node(String name, List<String> values) {
        this.name = name;
        this.values = values;
        this.parents = new ArrayList<>();
        this.cpt = new HashMap<>();
    }

    public String getName() {
        return name;
    }

    public List<String> getValues() {
        return values;
    }

    public List<Node> getParents() {
        return parents;
    }

    public void addParent(Node parent) {
        this.parents.add(parent);
    }

    public void setCpt(Map<List<String>, double[]> cpt) {
        // 自动补全缺失的父节点组合
        List<List<String>> allParentCombinations = generateAllParentCombinations();
        for (List<String> combination : allParentCombinations) {
            if (!this.cpt.containsKey(combination)) {
                double[] defaultProbs = new double[values.size()];
                Arrays.fill(defaultProbs, 1.0 / values.size()); // 均匀分布
                this.cpt.put(combination, defaultProbs);
            }
        }
        this.cpt.putAll(cpt); // 允许用户覆盖默认值
    }
    private List<List<String>> generateAllParentCombinations() {
        List<List<String>> combinations = Arrays.asList(Collections.emptyList());
        for (Node parent : parents) {
            List<List<String>> newCombinations = new ArrayList<>();
            for (String parentValue : parent.getValues()) {
                for (List<String> combo : combinations) {
                    List<String> newCombo = new ArrayList<>(combo);
                    newCombo.add(parentValue);
                    newCombinations.add(newCombo);
                }
            }
            combinations = newCombinations;
        }
        return combinations;
    }

    /**
     * 获取给定父节点值的情况下，节点取特定值的概率
     */
    public double getProbability(String value, List<String> parentValues) {
        if (parents.isEmpty()) {
            // 没有父节点，直接返回先验概率
            int idx = values.indexOf(value);
            return cpt.get(Collections.emptyList())[idx];
        }

        // 检查父节点值是否有效
        if (!cpt.containsKey(parentValues)) {
            throw new IllegalArgumentException("未定义的父节点值组合: " + parentValues);
        }

        // 查找对应的值的概率
        int idx = values.indexOf(value);
        return cpt.get(parentValues)[idx];
    }
}